import sys
import os
import math
import random
import numpy
import argparse
from scipy.io import wavfile


class audioDatastore:
	def __init__(self,path):
		self.path=path
		self.Files=[]

		for file in os.listdir(path):
			# check only text files
			if file.endswith('.wav'):
				self.Files.append(file)				
	def get_full_path(pair):
		parent=[]
		while  True:
			names.append(pair.name)
			if pair.path_index<0:
				break
			pair= self.Dirs[pair.path_index]
					
		return os.path.join(reversed(nsmes))
		
	def file(self,n):
		return os.path.join(self.path,self.Files[n])
		
def randi(maxi,n):
    return [  random.randrange(maxi) for i in range(n)]
	
def randf(n,m):
    return [  random.uniform(0,m) for i in range(n)]	

def augmentDataset(datasetloc,to_datasetloc="generated"):
	adsBkg = audioDatastore(datasetloc)
	# Known sample rate of the data set
	fs = 16e3 
	segmentDuration = 1
	segmentSamples = round(segmentDuration*fs)

	volumeRange = [math.log10(1e-4),math.log10(1)]

	numBkgSegments = 4000
	numBkgFiles = len(adsBkg.Files)
	numSegmentsPerFile = math.floor(numBkgSegments/numBkgFiles)
	
	fpTrain = os.path.join(to_datasetloc,"train","background")
	fpValidation = os.path.join(to_datasetloc,"validation","background")

	if not os.path.exists(fpTrain):

		# Create directories
		os.makedirs(fpTrain)
		os.makedirs(fpValidation)
		
		for backgroundFileIndex in range(len(adsBkg.Files)):
			fn=adsBkg.file(backgroundFileIndex)
			fn_base=adsBkg.Files[backgroundFileIndex][:-4]

			sample_rate,data= wavfile.read(fn)
			bkgFile = data/max(data)
			
			if bkgFile.ndim>1 and bkgFile.shape[0]!=1:
			    raise ValueError("only support 1 channel:(%s)"%bkgFile.shape)

			# Determine starting index of each segment
			segmentStart = randi(bkgFile.size-segmentSamples,numSegmentsPerFile);

			# Determine gain of each clip
			gain=[math.pow(10,((volumeRange[1]-volumeRange[0])*random.uniform(0,1) + volumeRange[0]))  for i in range(numSegmentsPerFile)]

			for segmentIdx in range(numSegmentsPerFile):

				# Isolate the randomly chosen segment of data.
				bkgSegment = bkgFile[segmentStart[segmentIdx]:segmentStart[segmentIdx]+segmentSamples-1];

				# Scale the segment by the specified gain.
				bkgSegment = bkgSegment*gain[segmentIdx];

				# Clip the audio between -1 and 1.
				bkgSegment = numpy.clip(bkgSegment,a_max=1,a_min=-1)

				# Create a file name.
				afn = fn_base + "_segment" + str(segmentIdx) + ".wav";

				# Randomly assign background segment to either the train or validation set.
				#  Assign 15% to validation
				if random.uniform(0,1) > 0.85:
					dirToWriteTo = fpValidation;
				else:
				# Assign 85% to train set.
					dirToWriteTo = fpTrain;

				# Write the audio to the file location.
				ffn = os.path.join(dirToWriteTo,afn);
				wavfile.write(ffn,sample_rate,bkgSegment)

			

			# Print progress
			print('Progress = %d (%%)\n'%round(100*backgroundFileIndex/len(adsBkg.Files)))



parser = argparse.ArgumentParser(description='augmrnt data set.')
parser.add_argument("--src", required=True)
parser.add_argument("--dst", required=True)

args = parser.parse_args()
augmentDataset(args.src,args.dst)
